//
// Created by myrica on 3/3/25.
//

#ifndef SCOPE_H
#define SCOPE_H

#include <map>
#include <ranges>
#include <string>
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"

#include "mlir/IR/Value.h"

// struct Scope {
// private:
//     std::map<std::string, mlir::Value *> variables;
//     std::unique_ptr<Scope> outer;
//
//     Scope(std::unique_ptr<Scope> __outer) : outer(std::move(__outer)) {
//     };
//
// public:
//     auto get(const std::string &name) {
//         if (const auto &res = variables.find(name); res != variables.end())
//             return res->second;
//         if (outer)
//             return outer->get(name);
//         return nullptr;
//     }
//
//     void set(const std::string &name, mlir::Value *value) {
//         variables.insert_or_assign(name, value);
//     }
//
//     // auto enterScope() {
//     //     return new Scope(this);
//     // }
//
//     // auto leaveScope() {
//     //     auto ptr = outer;
//     //     delete this;
//     //     return ptr;
//     // }
//     //
//
//     static std::unique_ptr<Scope> create(std::unique_ptr<Scope> outer = nullptr) {
//         return std::make_unique<Scope>(outer);
//     }
// };

struct Scope {
private:
    //    struct SymtableGlobal{};
public:
    using symtable_value_t = std::tuple<mlir::Value, mlir::Type, mlir::TypedAttr>;
    using symtable_func_t = std::pair<mlir::func::FuncOp, mlir::FunctionType>;

    // values[0] -> global
    std::vector<std::map<std::string, symtable_value_t> > values;
    std::vector<std::map<std::string, symtable_func_t> > funcs;

    Scope() {
        values.emplace_back();
        funcs.emplace_back();
    }

    symtable_value_t get(const std::string &key) {
        for (auto &value: std::ranges::reverse_view(values)) {
            if (auto res = value.find(key); res != value.end())
                return res->second;
        }
        return {nullptr, nullptr, nullptr};
    }

    symtable_func_t getFunc(const std::string &key) {
        if (auto res = funcs[0].find(key); res != funcs[0].end())
            return res->second;
        return {nullptr, nullptr};
    }

    void set(const std::string &key, const symtable_value_t &value) {
        assert(!values.empty());
        values.rbegin()->insert_or_assign(key, value);
    }

    void set(const std::string &key, const symtable_func_t &value) {
        assert(!funcs.empty());
        assert(isGlobal());
        funcs.rbegin()->insert_or_assign(key, value);
    }

    void enterScope() {
        values.emplace_back();
        funcs.emplace_back();
    }

    void leaveScope() {
        assert(!values.empty());
        values.pop_back();
        assert(!funcs.empty());
        funcs.pop_back();
    }

    [[nodiscard]] bool isGlobal() const {
        assert(values.size() == funcs.size());
        return values.size() == 1;
    }
};

#endif //SCOPE_H
